{
 "cells": [
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'"
   ],
   "id": "547091dfffce2d9f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 导入依赖库\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms, datasets\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ],
   "id": "aabe3fb0fdaccc3a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 图像预处理配置\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=(0.1307,), std=(0.3081,))  # MNIST专用归一化参数\n",
    "])\n",
    "\n",
    "# 反归一化器（用于还原图像显示）\n",
    "inv_transform = transforms.Compose([\n",
    "    transforms.Normalize(mean=[0.0], std=[1/0.3081]),\n",
    "    transforms.Normalize(mean=[-0.1307], std=[1.0])\n",
    "])"
   ],
   "id": "ce9381ec4d3e4112",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 训练参数\n",
    "batch_size = 64\n",
    "epochs = 5\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"使用设备: {device}\")"
   ],
   "id": "5c04a53c87cfcfd0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 创建图像显示窗口（2行5列，显示10个样本）\n",
    "fig_img, axs_img = plt.subplots(2, 5, figsize=(12, 6))\n",
    "axs_img = axs_img.flatten()  # 转为1D列表便于循环"
   ],
   "id": "a94d72b4a59faa38",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 创建准确率曲线窗口\n",
    "fig_acc, ax_acc = plt.subplots(figsize=(8, 5))"
   ],
   "id": "d138e369d6fc5ea1",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 加载数据\n",
    "def load_data():\n",
    "    train_dataset = datasets.MNIST(\n",
    "        root='../dataset/mnist/',\n",
    "        train=True,\n",
    "        download=True,\n",
    "        transform=train_transform\n",
    "    )\n",
    "    test_dataset = datasets.MNIST(\n",
    "        root='../dataset/mnist/',\n",
    "        train=False,\n",
    "        download=True,\n",
    "        transform=train_transform\n",
    "    )\n",
    "    return (\n",
    "        DataLoader(train_dataset, batch_size=batch_size, shuffle=True),\n",
    "        DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "    )"
   ],
   "id": "abc27bfdfc5ec085",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 定义模型\n",
    "class MNIST_CNN(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = torch.nn.Conv2d(1, 10, 5)\n",
    "        self.conv2 = torch.nn.Conv2d(10, 20, 5)\n",
    "        self.pool = torch.nn.MaxPool2d(2)\n",
    "        self.fc = torch.nn.Linear(320, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(batch_size, -1)\n",
    "        return self.fc(x)"
   ],
   "id": "d4fa975a1fdfb4dc",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 训练函数\n",
    "def train_epoch(model, train_loader, criterion, optimizer, epoch):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    for batch_idx, (inputs, targets) in enumerate(train_loader):\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, targets)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "        if batch_idx % 300 == 299:\n",
    "            print(f\"[Epoch {epoch+1}/{epochs}, Batch {batch_idx+1}] 平均损失: {running_loss/300:.3f}\")\n",
    "            running_loss = 0.0\n",
    "# 测试函数\n",
    "def test_model(model, test_loader):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for inputs, targets in test_loader:\n",
    "            inputs, targets = inputs.to(device), targets.to(device)\n",
    "            outputs = model(inputs)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += targets.size(0)\n",
    "            correct += (predicted == targets).sum().item()\n",
    "    acc = 100 * correct / total\n",
    "    print(f\"测试集准确率: {acc:.1f}% [{correct}/{total}]\\n\")\n",
    "    return acc"
   ],
   "id": "ed3f03e6f4982b12",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 图像显示函数（Jupyter适配版）\n",
    "def show_images_with_results(model, test_loader):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        # 获取测试集第一个批次\n",
    "        data_iter = iter(test_loader)\n",
    "        inputs, targets = next(data_iter)\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "        outputs = model(inputs)\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "\n",
    "        # 清空旧图\n",
    "        for ax in axs_img:\n",
    "            ax.clear()\n",
    "            ax.axis('off')\n",
    "\n",
    "        # 显示前10个样本\n",
    "        for i in range(10):\n",
    "            # 还原图像\n",
    "            img_tensor = inputs[i].cpu()\n",
    "            img_tensor = inv_transform(img_tensor)\n",
    "            img = transforms.ToPILImage()(img_tensor)\n",
    "\n",
    "            # 显示图像\n",
    "            axs_img[i].imshow(img, cmap='gray')\n",
    "\n",
    "            # 添加标签\n",
    "            true = targets[i].item()\n",
    "            pred = predicted[i].item()\n",
    "            color = 'green' if true == pred else 'red'\n",
    "            axs_img[i].set_title(f\"真实: {true}\\n预测: {pred}\", color=color, fontsize=10)\n",
    "\n",
    "        plt.tight_layout()\n",
    "        fig_img.suptitle('样本图像与识别结果', fontsize=14, y=1.02)\n",
    "        plt.show()  # Jupyter显示关键\n",
    "\n",
    "\n",
    "# 准确率曲线绘制函数\n",
    "def plot_accuracy(epochs, accuracies):\n",
    "    ax_acc.clear()\n",
    "    ax_acc.plot(epochs, accuracies, 'bo-', linewidth=2, markersize=6)\n",
    "    ax_acc.set_xlabel('训练轮次')\n",
    "    ax_acc.set_ylabel('准确率 (%)')\n",
    "    ax_acc.set_title('测试集准确率变化')\n",
    "    ax_acc.set_ylim(0, 100)\n",
    "    ax_acc.grid(True, linestyle='--', alpha=0.7)\n",
    "    plt.show()  # Jupyter显示关键"
   ],
   "id": "5334dced9b27d52d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 主函数\n",
    "def main():\n",
    "    train_loader, test_loader = load_data()\n",
    "    model = MNIST_CNN().to(device)\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)\n",
    "\n",
    "    # 训练前显示初始预测\n",
    "    print(\"===== 训练前：随机预测结果 =====\")\n",
    "    show_images_with_results(model, test_loader)\n",
    "\n",
    "    # 训练循环\n",
    "    epochs_list = []\n",
    "    acc_list = []\n",
    "    for epoch in range(epochs):\n",
    "        print(f\"\\n===== 训练轮次 {epoch+1}/{epochs} =====\")\n",
    "        train_epoch(model, train_loader, criterion, optimizer, epoch)\n",
    "        acc = test_model(model, test_loader)\n",
    "\n",
    "        epochs_list.append(epoch+1)\n",
    "        acc_list.append(acc)\n",
    "        plot_accuracy(epochs_list, acc_list)\n",
    "        show_images_with_results(model, test_loader)\n",
    "\n",
    "    print(\"训练完成！\")\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ],
   "id": "471113ac751cf6b",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
